# 图检索 Agent（Neo4j）

# graph_retriever.py
from neo4j import GraphDatabase
from typing import Dict, List, Any, Optional
import logging

class GraphRetriever:
    """图模式检索器 - 支持完整NL2SQL流程"""
    
    def __init__(self, neo4j_config):
        self.neo4j_driver = GraphDatabase.driver(
            neo4j_config['uri'], 
            auth=(neo4j_config['user'], neo4j_config['password'])
        )
        self.logger = logging.getLogger(__name__)
    
    def retrieve_schema_context(self, entities: Dict, database_name: str = None) -> Dict[str, Any]:
        """
        基于实体识别结果检索数据库模式
        
        Args:
            entities: 实体识别结果
            database_name: 目标数据库名称
            
        Returns:
            模式上下文信息，包含SQL生成所需的所有结构信息
        """
        self.logger.info("开始检索数据库模式上下文")
        
        context = {
            "database": database_name,
            # 表结构信息
            "tables": {},
            "actual_tables_found": [],
            # 关系信息
            "relationships": [],
            "foreign_keys": [],
            "primary_keys": [],
            # 数据类型信息
            "data_types": {},
            "column_details": {},
            # SQL生成支持
            "suggested_joins": [],
            "query_patterns": [],
            # 查询优化建议
            "query_insights": [],
            "optimization_hints": [],
            # 处理元数据
            "processing_time": None,
            "confidence": "medium"
        }
        
        start_time = time.time()
        
        try:
            # 获取实际表信息
            actual_tables = self._get_database_tables(database_name)
            context["actual_tables_found"] = actual_tables
            
            # 获取识别出的表的结构信息
            recognized_tables = [t["name"] for t in entities.get("tables", [])]
            tables_to_retrieve = [t for t in recognized_tables if t in actual_tables]
            
            if not tables_to_retrieve:
                # 如果没有精确匹配，使用所有实际表
                tables_to_retrieve = actual_tables[:3]  # 限制数量避免过多
            
            for table_name in tables_to_retrieve:
                table_info = self._get_table_details(database_name, table_name)
                if table_info:
                    context["tables"][table_name] = table_info
            
            # 获取关系信息
            if len(tables_to_retrieve) >= 1:
                relationships = self._get_relationships_for_tables(database_name, tables_to_retrieve)
                context["relationships"] = relationships
                
                # 提取外键和主键信息
                context["foreign_keys"] = self._extract_foreign_keys(relationships)
                context["primary_keys"] = self._extract_primary_keys(context["tables"])
            
            # 获取数据类型信息
            context["data_types"] = self._get_data_types_context(entities, database_name)
            context["column_details"] = self._get_column_details_context(context["tables"])
            
            # 生成连接建议
            context["suggested_joins"] = self._suggest_joins(context["relationships"], entities)
            
            # 识别查询模式
            context["query_patterns"] = self._identify_query_patterns(entities, context)
            
            # 生成查询洞察和优化提示
            context["query_insights"] = self._generate_query_insights(entities, context)
            context["optimization_hints"] = self._generate_optimization_hints(entities, context)
            
            # 计算置信度
            context["confidence"] = self._calculate_context_confidence(context)
            context["processing_time"] = time.time() - start_time
            
            self.logger.info(f"模式检索完成: {len(context['tables'])}表, {len(context['relationships'])}关系")
            return context
            
        except Exception as e:
            self.logger.error(f"模式检索失败: {e}")
            context["error"] = str(e)
            context["processing_time"] = time.time() - start_time
            return context
    
    def _get_database_tables(self, database_name: str) -> List[str]:
        """从图数据库获取指定数据库的所有表"""
        tables = []
        
        try:
            with self.neo4j_driver.session() as session:
                result = session.run("""
                    MATCH (t:Table {database: $db_name})
                    RETURN t.name as table_name
                    ORDER BY t.name
                """, db_name=database_name)
                
                for record in result:
                    tables.append(record["table_name"])
        except Exception as e:
            self.logger.warning(f"获取数据库 {database_name} 的表失败: {e}")
        
        return tables
    
    def _get_table_details(self, database_name: str, table_name: str) -> Dict:
        """获取表的详细信息"""
        with self.neo4j_driver.session() as session:
            # 获取表基本信息
            result = session.run("""
                MATCH (t:Table {name: $table_name, database: $db_name})
                RETURN t.column_count as column_count,
                       t.primary_key_count as pk_count,
                       t.foreign_key_count as fk_count
            """, table_name=table_name, db_name=database_name)
            
            table_info = result.single()
            if not table_info:
                return None
            
            # 获取列详细信息
            columns_result = session.run("""
                MATCH (t:Table {name: $table_name, database: $db_name})-[:HAS_COLUMN]->(c:Column)
                RETURN c.name as name, c.type as type, c.is_primary as is_primary,
                       c.is_unique as is_unique, c.is_foreign as is_foreign,
                       c.nullable as nullable
                ORDER BY c.is_primary DESC, c.name
            """, table_name=table_name, db_name=database_name)
            
            columns = []
            for record in columns_result:
                columns.append({
                    "name": record["name"],
                    "type": record["type"],
                    "is_primary": record["is_primary"],
                    "is_unique": record["is_unique"],
                    "is_foreign": record["is_foreign"],
                    "nullable": record["nullable"]
                })
            
            return {
                "name": table_name,
                "column_count": table_info["column_count"],
                "primary_key_count": table_info["pk_count"],
                "foreign_key_count": table_info["fk_count"],
                "columns": columns
            }
    
    def _get_relationships_for_tables(self, database_name: str, table_names: List[str]) -> List[Dict]:
        """获取表间关系"""
        relationships = []
        
        if len(table_names) < 1:
            return relationships
        
        with self.neo4j_driver.session() as session:
            # 查找表间所有关系
            result = session.run("""
                MATCH (c1:Column {database: $db_name})-[r:REFERENCES]->(c2:Column {database: $db_name})
                WHERE c1.table_name IN $table_names OR c2.table_name IN $table_names
                RETURN c1.table_name as source_table, c1.name as source_column,
                       c2.table_name as target_table, c2.name as target_column,
                       r.confidence as confidence, r.relationship_type as relationship_type,
                       r.description as description
                ORDER BY r.confidence DESC
            """, db_name=database_name, table_names=table_names)
            
            for record in result:
                relationship = {
                    "source_table": record["source_table"],
                    "source_column": record["source_column"],
                    "target_table": record["target_table"],
                    "target_column": record["target_column"],
                    "confidence": record["confidence"],
                    "type": record["relationship_type"],
                    "description": record["description"]
                }
                relationships.append(relationship)
        
        return relationships
    
    def _extract_foreign_keys(self, relationships: List[Dict]) -> List[Dict]:
        """从关系中提取外键信息"""
        foreign_keys = []
        
        for relationship in relationships:
            if relationship["confidence"] in ["high", "medium"]:
                foreign_key = {
                    "table": relationship["source_table"],
                    "column": relationship["source_column"],
                    "references_table": relationship["target_table"],
                    "references_column": relationship["target_column"]
                }
                foreign_keys.append(foreign_key)
        
        return foreign_keys
    
    def _extract_primary_keys(self, tables: Dict) -> List[Dict]:
        """从表结构中提取主键信息"""
        primary_keys = []
        
        for table_name, table_info in tables.items():
            for column in table_info.get("columns", []):
                if column.get("is_primary"):
                    primary_key = {
                        "table": table_name,
                        "column": column["name"],
                        "type": column["type"]
                    }
                    primary_keys.append(primary_key)
        
        return primary_keys
    
    def _get_data_types_context(self, entities: Dict, database_name: str) -> Dict:
        """获取数据类型上下文"""
        data_types = {}
        
        for column in entities.get("columns", []):
            table_name = column.get("table")
            column_name = column["name"]
            
            if table_name:
                column_type = self._get_column_type(database_name, table_name, column_name)
                if column_type:
                    data_types[f"{table_name}.{column_name}"] = column_type
        
        return data_types
    
    def _get_column_details_context(self, tables: Dict) -> Dict:
        """获取列详情上下文"""
        column_details = {}
        
        for table_name, table_info in tables.items():
            for column in table_info.get("columns", []):
                key = f"{table_name}.{column['name']}"
                column_details[key] = {
                    "table": table_name,
                    "column": column["name"],
                    "type": column["type"],
                    "is_primary": column["is_primary"],
                    "is_foreign": column["is_foreign"],
                    "nullable": column["nullable"]
                }
        
        return column_details
    
    def _get_column_type(self, database_name: str, table_name: str, column_name: str) -> str:
        """获取列的数据类型"""
        with self.neo4j_driver.session() as session:
            result = session.run("""
                MATCH (c:Column {name: $column_name, table_name: $table_name, database: $db_name})
                RETURN c.type as type
            """, column_name=column_name, table_name=table_name, db_name=database_name)
            
            record = result.single()
            return record["type"] if record else None
    
    def _suggest_joins(self, relationships: List[Dict], entities: Dict) -> List[Dict]:
        """生成连接建议"""
        suggestions = []
        
        tables = [t["name"] for t in entities.get("tables", [])]
        
        if len(tables) < 2:
            return suggestions
        
        for relationship in relationships:
            if (relationship["source_table"] in tables and 
                relationship["target_table"] in tables):
                suggestion = {
                    "left_table": relationship["source_table"],
                    "left_column": relationship["source_column"],
                    "right_table": relationship["target_table"],
                    "right_column": relationship["target_column"],
                    "join_type": "INNER JOIN",
                    "confidence": relationship["confidence"]
                }
                suggestions.append(suggestion)
        
        return suggestions
    
    def _identify_query_patterns(self, entities: Dict, context: Dict) -> List[str]:
        """识别查询模式"""
        patterns = []
        
        # 简单查询模式
        if (len(entities.get("tables", [])) == 1 and 
            not entities.get("aggregations") and 
            not entities.get("time_conditions")):
            patterns.append("simple_select")
        
        # 聚合查询模式
        if entities.get("aggregations"):
            patterns.append("aggregation_query")
        
        # 时间范围查询模式
        if entities.get("time_conditions"):
            patterns.append("temporal_query")
        
        # 多表连接模式
        if len(context.get("suggested_joins", [])) > 0:
            patterns.append("multi_table_join")
        
        # 条件过滤模式
        if entities.get("conditions") or entities.get("values"):
            patterns.append("filtered_query")
        
        return patterns
    
    def _generate_query_insights(self, entities: Dict, context: Dict) -> List[str]:
        """生成查询洞察"""
        insights = []
        
        tables = entities.get("tables", [])
        aggregations = entities.get("aggregations", [])
        time_conditions = entities.get("time_conditions", [])
        relationships = context.get("relationships", [])
        
        if len(tables) == 1:
            insights.append("单表查询，结构简单")
        elif len(tables) > 1:
            insights.append(f"多表查询({len(tables)}个表)，可能需要表连接")
        
        if aggregations:
            agg_types = [agg["type"] for agg in aggregations]
            insights.append(f"包含聚合操作: {', '.join(agg_types)}")
        
        if time_conditions:
            time_texts = [tc["matched_text"] for tc in time_conditions]
            insights.append(f"包含时间条件: {', '.join(time_texts)}")
        
        if relationships:
            insights.append(f"发现{len(relationships)}个表关系")
        
        # 基于查询模式的洞察
        patterns = context.get("query_patterns", [])
        if "aggregation_query" in patterns:
            insights.append("建议使用GROUP BY进行分组统计")
        if "temporal_query" in patterns:
            insights.append("时间条件查询，注意日期格式")
        
        return insights
    
    def _generate_optimization_hints(self, entities: Dict, context: Dict) -> List[str]:
        """生成优化提示"""
        hints = []
        
        tables = context.get("tables", {})
        
        # 检查是否有主键用于优化
        for table_name, table_info in tables.items():
            if table_info.get("primary_key_count", 0) > 0:
                hints.append(f"表 {table_name} 有主键，查询效率较高")
        
        # 检查大数据量表
        large_tables = [name for name, info in tables.items() 
                       if info.get("column_count", 0) > 10]
        if large_tables:
            hints.append(f"表 {', '.join(large_tables)} 列数较多，建议指定需要的列")
        
        # 多表连接优化提示
        if len(context.get("suggested_joins", [])) > 2:
            hints.append("多表连接查询，考虑查询性能优化")
        
        return hints
    
    def _calculate_context_confidence(self, context: Dict) -> str:
        """计算上下文置信度"""
        tables_found = len(context.get("actual_tables_found", []))
        relationships_found = len(context.get("relationships", []))
        
        if tables_found == 0:
            return "low"
        elif relationships_found > 0 and tables_found >= 1:
            return "high"
        elif tables_found >= 1:
            return "medium"
        else:
            return "low"
    
    def close(self):
        """关闭Neo4j连接"""
        if self.neo4j_driver:
            self.neo4j_driver.close()

# 添加time模块导入
import time